package com.demo.spring.rag2;

import java.time.Instant;
import java.util.ArrayList;
import java.util.List;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

@Configuration
public class CommonConfig {

	@Autowired
	private OpenAiChatModel model;
	@Autowired
	private ChatMemoryStore store;
	@Autowired
	private EmbeddingModel embeddingModel;
	
	// 构建会话记忆体
	@Bean
	public ChatMemoryProvider chatMemoryProvider() {
		ChatMemoryProvider chatMemoryProvider = new ChatMemoryProvider() {

			public ChatMemory get(Object memoryId) {
				MessageWindowChatMemory memory = MessageWindowChatMemory.builder()
						.id(memoryId)
						.chatMemoryStore(store)
						.maxMessages(20)
						.build();
				return memory;
			}
			
		};
		return chatMemoryProvider;
	}
	
	// 构建向量数据库
	@Bean
	@Primary
	public EmbeddingStore store() {
		// 1、加载文件进内存
		List<Document> documents = FileSystemDocumentLoader.loadDocuments("D:/doc", new ApachePdfBoxDocumentParser());
		// 2、构建向量数据库操作对象
		InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
		// 3、构建文档分割器
		DocumentSplitter ds = DocumentSplitters.recursive(300, 100);
		List<TextSegment> chunks = ds.splitAll(documents);
		// 定义元数据提取逻辑
		for (TextSegment chunk : chunks) {
		    // 1. 创建新Metadata对象
		    Metadata metadata = new Metadata();
		    
		    // 2. 添加原始文档的元数据（如文件名）
		    metadata.put("source", documents.get(0).metadata().getString("file_name")); 
		    
		    // 3. 添加自定义元数据（如时间戳）
		    metadata.put("timestamp", Instant.now().toString());
		    
		    // 4. 重建TextSegment对象
		    TextSegment enrichedChunk = new TextSegment(
		        chunk.text(),  // 保留原始文本
		        metadata      // 注入新元数据
		    );
		    
		    // 替换原始chunk
		    chunks.set(chunks.indexOf(chunk), enrichedChunk);
		}
		
		// 3. 手动分批处理（每批10个chunk）
		int batchSize = 10; // 严格遵循API限制
		for (int i = 0; i < chunks.size(); i += batchSize) {
		    List<TextSegment> batch = chunks.subList(i, Math.min(i + batchSize, chunks.size()));
		    
		    // 生成当前批次的向量
		    List<Embedding> embeddings = new ArrayList<Embedding>();
		    for (TextSegment chunk : batch) {
		        embeddings.add(embeddingModel.embed(chunk.text()).content());
		    }
		    
		    // 存储向量和原始文本
		    store.addAll(embeddings, batch);
		}
		return store;
	}
	
	// 构建向量数据库检索对象
	@Bean
	public ContentRetriever contentRetriever(EmbeddingStore store) {
		return EmbeddingStoreContentRetriever.builder()
				.embeddingStore(store)
				.minScore(0.5)
				.maxResults(3)
				.embeddingModel(embeddingModel)
				.build();
	}
	
}
